import pandas as pd
import numpy as np
import pymysql
from datetime import datetime, timedelta
import uuid
import warnings
warnings.filterwarnings('ignore')

# 从 database.py 导入数据库配置
from database import DB_HOST, DB_USER, DB_PASSWORD, DB_NAME, DB_PORT

# 使用从 database.py 导入的配置
db_config = {
    'host': DB_HOST,
    'user': DB_USER,
    'password': DB_PASSWORD,
    'database': DB_NAME,
    'charset': 'utf8mb4',
    'port': int(DB_PORT) 
}


def load_recent_data(days=90):
    """加载最近N天的数据（用于任务1）"""
    connection = pymysql.connect(**db_config)

    cutoff_date = (datetime.now() - timedelta(days=days)).strftime('%Y-%m-%d')

    query = f"""
    SELECT 
        EFAlarmID, FenceName, FenceLocation as location,
        StartTime, EndTime, FenceState
    FROM ef_alarm 
    WHERE StartTime >= '{cutoff_date}' 
    AND StartTime IS NOT NULL
    ORDER BY StartTime
    """

    try:
        df = pd.read_sql(query, connection)
        print(f"加载到最近{days}天内 {len(df)} 条报警记录")
    except Exception as e:
        print(f"数据库查询错误: {e}")
        return pd.DataFrame()
    finally:
        connection.close()

    return process_time_data(df)


def load_all_historical_data():
    """加载所有历史数据（用于任务2）"""
    connection = pymysql.connect(**db_config)

    query = """
    SELECT 
        EFAlarmID, FenceName, FenceLocation as location,
        StartTime, EndTime, FenceState
    FROM ef_alarm 
    WHERE StartTime IS NOT NULL
    ORDER BY StartTime
    """

    try:
        df = pd.read_sql(query, connection)
        print(f"加载到所有历史数据 {len(df)} 条报警记录")
    except Exception as e:
        print(f"数据库查询错误: {e}")
        return pd.DataFrame()
    finally:
        connection.close()

    return process_time_data(df)


def process_time_data(df):
    """处理时间数据"""
    if df.empty:
        return df

    df['StartTime'] = pd.to_datetime(df['StartTime'], errors='coerce')
    df['EndTime'] = pd.to_datetime(df['EndTime'], errors='coerce')
    df = df[df['StartTime'].notna()]

    # 提取小时信息
    df['hour'] = df['StartTime'].dt.hour

    return df


def calculate_time_decay_risk(df):
    """基于时间衰减计算区域风险等级 - 修复版本"""
    if df.empty:
        return pd.DataFrame()

    risk_results = []
    current_time = datetime.now()

    for location in df['location'].unique():
        location_data = df[df['location'] == location]
        device_name = location_data['FenceName'].iloc[0] if len(location_data) > 0 else '未知设备'

        # 生成唯一记录ID
        record_id = str(uuid.uuid4())

        # 计算时间衰减权重总分
        total_time_weight = 0
        latest_alarm = None
        alarm_count = len(location_data)

        print(f"区域 '{location}' 有 {alarm_count} 条报警记录")

        for _, row in location_data.iterrows():
            days_ago = (current_time - row['StartTime']).days

            # 改进的时间衰减函数：使用更平缓的衰减
            if days_ago <= 365:
                # 使用指数衰减而不是线性衰减，保留更多权重
                time_weight = 0.95 ** days_ago  # 指数衰减，90天后约剩0.01
            else:
                time_weight = 0

            total_time_weight += time_weight

            # 记录最新报警时间
            if latest_alarm is None or row['StartTime'] > latest_alarm:
                latest_alarm = row['StartTime']

        print(f"区域 '{location}' 总时间权重: {total_time_weight:.4f}")

        # 改进的风险分数计算
        if alarm_count > 0:
            # 基础风险：只要有报警就有基础风险
            base_risk = 0.3

            # 频率风险：报警次数的影响（使用对数函数避免过大）
            frequency_risk = min(np.log1p(alarm_count) / 3, 0.4)  # 最大0.4

            # 时间风险：时间衰减权重的影响
            time_risk = min(total_time_weight / 2, 0.3)  # 最大0.3

            risk_score = base_risk + frequency_risk + time_risk
        else:
            risk_score = 0

        # 确保风险分数在0-1之间
        risk_score = min(risk_score, 1.0)

        print(f"区域 '{location}' 最终风险分数: {risk_score:.4f}")

        # 调整风险等级阈值，更适合低频系统
        if risk_score >= 0.9:
            risk_level = "超高风险"
        elif risk_score >= 0.7:
            risk_level = "高风险"
        elif risk_score >= 0.5:
            risk_level = "中风险"
        elif risk_score >= 0.3:
            risk_level = "低风险"
        else:
            risk_level = "无风险"

        risk_results.append({
            'record_id': record_id,
            'device_name': device_name,
            'location': location,
            'risk_score': round(risk_score, 4),
            'risk_level': risk_level,
            'alarm_count': alarm_count,
            'total_time_weight': round(total_time_weight, 4),
            'last_alarm_date': latest_alarm,
            'update_time': current_time
        })

    return pd.DataFrame(risk_results)


def save_risk_level_to_db(risk_df, table_name='region_risk_level'):
    """保存风险等级到数据库"""
    if risk_df.empty:
        print("没有风险等级数据需要保存")
        return

    connection = pymysql.connect(**db_config)

    try:
        create_table_sql = f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            id INT AUTO_INCREMENT PRIMARY KEY,
            record_id VARCHAR(36) NOT NULL,
            device_name VARCHAR(100),
            location VARCHAR(100),
            risk_score DECIMAL(8,4),
            risk_level VARCHAR(20),
            alarm_count INT,
            total_time_weight DECIMAL(8,4),
            last_alarm_date DATETIME,
            update_time DATETIME,
            created_time DATETIME DEFAULT CURRENT_TIMESTAMP,
            UNIQUE KEY unique_record (record_id),
            INDEX idx_location (location),
            INDEX idx_update_time (update_time)
        )
        """

        with connection.cursor() as cursor:
            cursor.execute(create_table_sql)

        # 使用INSERT INTO添加新记录，不覆盖旧记录
        insert_count = 0
        for _, row in risk_df.iterrows():
            query = f"""
            INSERT INTO {table_name} 
            (record_id, device_name, location, risk_score, risk_level, alarm_count, total_time_weight, last_alarm_date, update_time)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            """

            with connection.cursor() as cursor:
                cursor.execute(query, (
                    row['record_id'], row['device_name'], row['location'],
                    row['risk_score'], row['risk_level'], row['alarm_count'],
                    row['total_time_weight'], row['last_alarm_date'], row['update_time']
                ))
                insert_count += 1

        connection.commit()
        print(f"成功保存 {insert_count} 条风险等级记录到 {table_name}")

    except Exception as e:
        print(f"保存风险等级时出错: {e}")
        connection.rollback()
    finally:
        connection.close()


def calculate_hourly_frequency_risk(df):
    """基于小时频率计算24小时风险分布"""
    if df.empty:
        return pd.DataFrame()

    hourly_results = []

    for location in df['location'].unique():
        location_data = df[df['location'] == location]
        device_name = location_data['FenceName'].iloc[0] if len(location_data) > 0 else '未知设备'

        # 生成唯一记录ID
        record_id = str(uuid.uuid4())

        # 统计每个小时的报警次数
        hour_counts = location_data['hour'].value_counts().sort_index()

        # 找到最大次数用于归一化
        max_count = hour_counts.max() if len(hour_counts) > 0 else 1

        # 创建24小时风险记录
        record = {
            'record_id': record_id,
            'device_name': device_name,
            'location': location,
            'update_time': datetime.now()
        }

        # 为每个小时计算风险值
        for hour in range(24):
            count = hour_counts.get(hour, 0)
            risk_value = count / max_count  # 归一化到0-1
            record[f'h{hour:02d}_risk'] = round(risk_value, 4)

        hourly_results.append(record)

    return pd.DataFrame(hourly_results)


def save_hourly_risk_to_db(hourly_df, table_name='hourly_risk_distribution'):
    """保存24小时风险分布到数据库"""
    if hourly_df.empty:
        print("没有小时风险数据需要保存")
        return

    connection = pymysql.connect(**db_config)

    try:
        # 动态创建表（包含24个小时字段）
        columns_sql = []
        for hour in range(24):
            columns_sql.append(f"h{hour:02d}_risk DECIMAL(6,4) DEFAULT 0")

        create_table_sql = f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            id INT AUTO_INCREMENT PRIMARY KEY,
            record_id VARCHAR(36) NOT NULL,
            device_name VARCHAR(100),
            location VARCHAR(100),
            {', '.join(columns_sql)},
            update_time DATETIME,
            created_time DATETIME DEFAULT CURRENT_TIMESTAMP,
            UNIQUE KEY unique_record (record_id),
            INDEX idx_location (location),
            INDEX idx_update_time (update_time)
        )
        """

        with connection.cursor() as cursor:
            cursor.execute(create_table_sql)

        # 准备插入语句
        hour_columns = [f"h{hour:02d}_risk" for hour in range(24)]
        columns = ['record_id', 'device_name', 'location'] + hour_columns + ['update_time']
        placeholders = ['%s'] * len(columns)

        insert_sql = f"""
        INSERT INTO {table_name} ({', '.join(columns)})
        VALUES ({', '.join(placeholders)})
        """

        # 插入数据
        insert_count = 0
        for _, row in hourly_df.iterrows():
            values = [row['record_id'], row['device_name'], row['location']]
            for hour in range(24):
                values.append(row.get(f'h{hour:02d}_risk', 0.0))
            values.append(row['update_time'])

            with connection.cursor() as cursor:
                cursor.execute(insert_sql, values)
                insert_count += 1

        connection.commit()
        print(f"成功保存 {insert_count} 条小时风险分布记录到 {table_name}")

    except Exception as e:
        print(f"保存小时风险分布时出错: {e}")
        connection.rollback()
    finally:
        connection.close()


def main():
    """主执行流程"""
    print("开始红外电子围栏风险评估...")
    print("=" * 60)

    # 任务1：基于90天时间衰减的风险等级
    print("\n任务1: 计算180天时间衰减风险等级...")
    recent_df = load_recent_data(days=180)

    if not recent_df.empty:
        risk_level_df = calculate_time_decay_risk(recent_df)

        if not risk_level_df.empty:
            print("风险等级统计:")
            for level in ["超高风险", "高风险", "中风险", "低风险", "无风险"]:
                count = len(risk_level_df[risk_level_df['risk_level'] == level])
                if count > 0:
                    print(f"  {level}: {count} 个区域")

            # 显示高风险区域示例
            high_risk = risk_level_df[risk_level_df['risk_level'].isin(["超高风险", "高风险"])]
            if not high_risk.empty:
                print("\n高风险区域示例:")
                for _, row in high_risk.head(3).iterrows():
                    print(f"  区域: {row['location']}, 设备: {row['device_name']}, 等级: {row['risk_level']}")

            save_risk_level_to_db(risk_level_df)

    # 任务2：基于所有历史数据的小时频率风险分布
    print("\n任务2: 计算24小时频率风险分布...")
    all_df = load_all_historical_data()

    if not all_df.empty:
        hourly_risk_df = calculate_hourly_frequency_risk(all_df)

        if not hourly_risk_df.empty:
            # 显示某个区域的小时风险示例
            sample_location = hourly_risk_df['location'].iloc[0]
            sample_data = hourly_risk_df[hourly_risk_df['location'] == sample_location].iloc[0]

            print(f"\n区域 '{sample_location}' 的小时风险示例:")
            peak_hours = []
            for hour in range(24):
                risk_key = f'h{hour:02d}_risk'
                risk_value = sample_data.get(risk_key, 0)
                if risk_value > 0:
                    peak_hours.append((hour, risk_value))

            # 显示风险最高的3个时段
            peak_hours.sort(key=lambda x: x[1], reverse=True)
            for hour, risk in peak_hours[:3]:
                print(f"  {hour:02d}:00 - 风险值: {risk:.4f}")

            save_hourly_risk_to_db(hourly_risk_df)

    print("\n" + "=" * 60)
    print("风险评估完成!")

    return risk_level_df if 'risk_level_df' in locals() else None, hourly_risk_df if 'hourly_risk_df' in locals() else None


# 执行
if __name__ == "__main__":
    risk_level, hourly_risk = main()